import config
import copy
import pickle

from tabulate import tabulate


def SortPartitionsByTraceId(span_partitions):
    for ep, part in span_partitions.items():
        part.sort(key=lambda x: x.trace_id)


def SortPartitionsByTime(span_partitions):
    for ep, part in span_partitions.items():
        part.sort(key=lambda x: (x.start_mus, x.start_mus + x.duration_mus))


# 获取 out_ep 列表，其实也就是全体服务列表。
# 结果按照访问 out_ep 的 start 升序排序。
def GetOutEpsInOrder(out_span_partitions):
    eps = []
    for ep, spans in out_span_partitions.items():
        assert len(spans) > 0
        eps.append((ep, spans[0].start_mus))
    eps.sort(key=lambda x: x[1])
    return [x[0] for x in eps]


# 利用 trace_id 确定实际 mapping（GroundTruth）。
# 如果没有 trace_id，返回空。
def GetGroundTruth(in_span_partitions, out_span_partitions):
    assert len(in_span_partitions) == 1
    _, in_spans = list(in_span_partitions.items())[0]
    true_assignments = {ep: {} for ep in out_span_partitions.keys()}
    for in_span in in_spans:
        for out_ep in out_span_partitions.keys():
            for span in out_span_partitions[out_ep]:
                if span.trace_id == in_span.trace_id:
                    true_assignments[out_ep][in_span.GetId()] = span.GetId()
                    break
    return true_assignments


def GetGroundTruthMultiplex(in_span_partitions, out_span_partitions):
    final_true_assignments = {}
    for in_ep, in_span_partition in in_span_partitions.items():
        partial_in_span_partition = {in_ep: in_span_partition}
        partial_true_assignments = GetGroundTruth(partial_in_span_partition, out_span_partitions)

        for out_ep, assignments in partial_true_assignments.items():
            if out_ep not in final_true_assignments:
                final_true_assignments[out_ep] = {}
            final_true_assignments[out_ep].update(assignments)
    return final_true_assignments


def _AccuracyForSpan(pred_assignments, true_assignments, in_span_id):
    correct = True
    for ep in true_assignments.keys():
        if isinstance(pred_assignments[ep][in_span_id], list):
            if len(pred_assignments[ep][in_span_id]) > 1:
                correct = False
            else:
                pred_assignments[ep][in_span_id] = pred_assignments[ep][in_span_id][0]
        correct = correct and (
                pred_assignments[ep][in_span_id]
                == true_assignments[ep][in_span_id]
        )
    return int(correct)


def _TopKAccuracyForSpan(pred_topk_assignments, true_assignments, in_span_id):
    ep0 = list(true_assignments.keys())[0]
    correct = False
    for i in range(len(pred_topk_assignments[ep0][in_span_id])):
        correct = True
        for ep in true_assignments.keys():
            correct = correct and (
                    pred_topk_assignments[ep][in_span_id][i]
                    == true_assignments[ep][in_span_id]
            )
        if correct:
            break
    return int(correct)


def accuracyForService(pred_assignments, true_assignments, in_span_partitions) -> (int, int):
    assert len(in_span_partitions) == 1
    _, in_spans = list(in_span_partitions.items())[0]
    valid_cnt = 0  # 有效样本计数
    correct_cnt = 0  # 正确样本计数

    for in_span in in_spans:
        # 检查 in_span 是否为有效样本
        is_valid = True
        for ep in true_assignments.keys():
            if ep not in pred_assignments or \
                    in_span.GetId() not in true_assignments[ep] or \
                    in_span.GetId() not in pred_assignments[ep]:
                is_valid = False
                break

        if not is_valid:
            continue  # 跳过无效样本

        valid_cnt += 1
        correct = True
        for ep in true_assignments.keys():
            if isinstance(pred_assignments[ep][in_span.GetId()], list):
                if len(pred_assignments[ep][in_span.GetId()]) > 1:
                    correct = False  # fixme 为什么单个 parent span 存在多个 child span，这样会被判定 False？
                    # 我认为应该是用 child span 查找唯一的 parent span，然后对比。https://github.com/Sachin-A/TraceWeaver/issues/3
                else:
                    pred_assignments[ep][in_span.GetId()] = pred_assignments[ep][in_span.GetId()][0]
            correct = correct and (
                    pred_assignments[ep][in_span.GetId()]
                    == true_assignments[ep][in_span.GetId()]
            )
        correct_cnt += int(correct)

    print(f"AccuracyForService = {correct_cnt} / {valid_cnt}")
    return correct_cnt, valid_cnt


def AccuracyForService(pred_assignments, true_assignments, in_span_partitions):
    assert len(in_span_partitions) == 1

    correct_cnt, valid_cnt = accuracyForService(pred_assignments, true_assignments, in_span_partitions)
    if valid_cnt == 0:
        return 0.0  # 避免除零错误
    return correct_cnt / valid_cnt


def AccuracyForServiceMultiplex(pred_assignments, true_assignments, in_span_partitions):
    final_correct_cnt, final_valid_cnt = 0, 0
    for in_ep, in_span_partition in in_span_partitions.items():
        partial_in_span_partition = {in_ep: in_span_partition}
        partial_correct_cnt, partial_valid_cnt = accuracyForService(
            pred_assignments, true_assignments, partial_in_span_partition
        )
        final_correct_cnt += partial_correct_cnt
        final_valid_cnt += partial_valid_cnt
    if final_valid_cnt == 0:
        return 0.0  # 避免除零错误
    return float(final_correct_cnt) / float(final_valid_cnt)


def topKAccuracyForService(pred_topk_assignments, true_assignments, in_span_partitions):
    assert len(in_span_partitions) == 1
    _, in_spans = list(in_span_partitions.items())[0]
    valid_cnt = 0  # 有效样本计数
    correct_cnt = 0  # 正确样本计数

    ep0 = list(true_assignments.keys())[0]
    for in_span in in_spans:
        # 检查 in_span 是否为有效样本
        is_valid = True
        for ep in true_assignments.keys():
            if in_span.GetId() not in true_assignments[ep] or in_span.GetId() not in pred_topk_assignments[ep]:
                is_valid = False
                break

        if not is_valid:
            continue  # 跳过无效样本

        valid_cnt += 1
        correct = False
        for i in range(len(pred_topk_assignments[ep0][in_span.GetId()])):
            correct = True
            for ep in true_assignments.keys():
                correct = correct and (
                        pred_topk_assignments[ep][in_span.GetId()][i]
                        == true_assignments[ep][in_span.GetId()]
                )
            if correct:
                break
        correct_cnt += int(correct)

    return float(correct_cnt), valid_cnt


def TopKAccuracyForService(pred_topk_assignments, true_assignments, in_span_partitions):
    return 0  # 懒得改了，反正也不用 topk。

    assert len(in_span_partitions) == 1

    correct_cnt, valid_cnt = topKAccuracyForService(pred_topk_assignments, true_assignments, in_span_partitions)
    if valid_cnt == 0:
        return 0.0  # 避免除零错误
    return correct_cnt / valid_cnt


def TopKAccuracyForServiceMultiplex(pred_topk_assignments, true_assignments, in_span_partitions):
    return 0  # 懒得改了，反正也不用 topk。

    final_correct_cnt, final_valid_cnt = 0, 0
    for in_ep, in_span_partition in in_span_partitions.items():
        partial_in_span_partition = {in_ep: in_span_partition}
        partial_correct_cnt, partial_valid_cnt = topKAccuracyForService(
            pred_topk_assignments, true_assignments, partial_in_span_partition
        )
        final_correct_cnt += partial_correct_cnt
        final_valid_cnt += partial_valid_cnt
    if final_valid_cnt == 0:
        return 0.0  # 避免除零错误
    return float(final_correct_cnt) / float(final_valid_cnt)


def AccuracyEndToEnd(pred_assignments_by_process, true_assignments_by_process, in_spans_by_process, dont_want):
    """
    快照参考 /root/Pictures/Work/adjust_acc/in_spans结构.png

    :param pred_assignments_by_process: in_ep => out_ep => in_span => out_span
    :param true_assignments_by_process: in_ep => out_ep => in_span => out_span
    :param in_spans_by_process: in_ep => in_span
    :param dont_want: directly return 0
    :return:
    """

    if dont_want:
        return {}, 0

    processes = true_assignments_by_process.keys()
    processes = list(set(processes) - set(config.otel_block_list))

    trace_correctness = {}
    for process in processes:
        for in_span in in_spans_by_process[process]:
            if in_span.trace_id not in trace_correctness:
                trace_correctness[in_span.trace_id] = True
            true_assignments = true_assignments_by_process[process]
            pred_assignments = pred_assignments_by_process[process]
            for out_ep in true_assignments.keys():
                if in_span.GetId() not in true_assignments[out_ep] or \
                        in_span.GetId() not in pred_assignments[out_ep]:
                    # try del
                    if in_span.trace_id in trace_correctness:
                        del trace_correctness[in_span.trace_id]
                    continue
                if (
                        true_assignments[out_ep][in_span.GetId()]
                        != pred_assignments[out_ep][in_span.GetId()]
                ):
                    trace_correctness[in_span.trace_id] = False
    correct = sum(trace_correctness[tid] for tid in trace_correctness)
    print(f"AccuracyEndToEnd = {correct} / {len(trace_correctness)}")
    if len(trace_correctness) == 0:
        return {}, 0.0
    return trace_correctness, float(correct) / len(trace_correctness)


def TopKAccuracyEndToEnd(pred_topk_assignments_by_process, true_assignments_by_process, in_spans_by_process,
                         dont_want):
    return {}, 0  # 懒得改了，反正也不用 topk。

    if dont_want:
        return {}, 0

    processes = true_assignments_by_process.keys()
    trace_acc = {}
    for i, process in enumerate(processes):
        true_assignments = true_assignments_by_process[process]
        pred_topk_assignments = pred_topk_assignments_by_process[process]
        ep0 = list(true_assignments.keys())[0]
        for x, in_span in enumerate(in_spans_by_process[process]):
            if i != 0 and trace_acc[in_span.trace_id] == False:
                continue
            if len(pred_topk_assignments[ep0][in_span.GetId()]) < 1:
                trace_acc[in_span.trace_id] = False
                continue
            for j in range(len(pred_topk_assignments[ep0][in_span.GetId()])):
                trace_acc[in_span.trace_id] = True
                for ep in true_assignments.keys():
                    if in_span.GetId() not in true_assignments[ep]:
                        # trace_acc[in_span.trace_id] = False
                        del trace_acc[in_span.trace_id]
                        continue
                    if (
                            true_assignments[ep][in_span.GetId()]
                            != pred_topk_assignments[ep][in_span.GetId()][j]
                    ):
                        trace_acc[in_span.trace_id] = False
                if trace_acc.get(in_span.trace_id, True):
                    break
    correct = sum(trace_acc[tid] for tid in trace_acc)
    if len(trace_acc) == 0:
        return {}, 0.0
    return trace_acc, float(correct) / len(trace_acc)


# 复用了 e2e_traces，上面的计算结果
def _BinAccuracyByServiceTimes(method, PLOTS_DIR, dont_want):
    if dont_want:
        return 0

    for j in range(4, 5):

        query_latency = {}

        with open(PLOTS_DIR + "e2e_" + str((j + 1) * 25) + ".pickle", 'rb') as afile:
            e2e_traces = pickle.load(afile)

        true_traces = e2e_traces[method][0]
        pred_traces = e2e_traces[method][1]

        all_traces = []

        for trace in true_traces.items():
            true_trace = true_traces[trace[0]]
            duration = true_trace[1].start_mus - (true_trace[0].start_mus + true_trace[0].duration_mus)
            if pred_traces[trace[0]][1]:
                correct = true_trace[1].sid == pred_traces[trace[0]][1].sid
            else:
                correct = False
            all_traces.append((duration, correct, 1))

        all_traces.sort()
        for i in range(1, len(all_traces)):
            _, c, n = all_traces[i - 1]
            t0, c0, n0 = all_traces[i]
            all_traces[i] = (t0, c + c0, n + n0)
        nbins = 10
        prev_c, prev_n = 0, 0
        accuracy = []
        for b in range(nbins):
            d, c, n = all_traces[int((len(all_traces) * (b + 1)) / nbins - 1)]
            c, n = c - prev_c, n - prev_n
            prev_c, prev_n = prev_c + c, prev_n + n
            percentile = (b + 1) * 100 / nbins
            acc = c / n
            accuracy.append((percentile, acc, d / 1000.0))
        return accuracy


def BinAccuracyByResponseTimes(trace_correctness, all_spans, dont_want):
    if dont_want:
        return 0

    all_traces = []
    for _, span in all_spans.items():
        if span.IsRoot() and span.trace_id in trace_correctness:
            correct = trace_correctness[span.trace_id]
            all_traces.append((span.duration_mus, span.trace_id, correct, 1))
    all_traces.sort()
    # accumulate
    for i in range(1, len(all_traces)):
        _, _, c, n = all_traces[i - 1]
        t0, s0, c0, n0 = all_traces[i]
        all_traces[i] = (t0, s0, c + c0, n + n0)
    nbins = 10
    prev_c, prev_n = 0, 0
    accuracy = []
    if not all_traces:
        return accuracy
    for b in range(nbins):
        d, _, c, n = all_traces[int((len(all_traces) * (b + 1)) / nbins - 1)]
        c, n = c - prev_c, n - prev_n
        prev_c, prev_n = prev_c + c, prev_n + n
        percentile = (b + 1) * 100 / nbins
        if n == 0:
            acc = 0
        else:
            acc = c / n
        if config.VERBOSE:
            print(
                "Accuracy of %d-percentile bin: %.3f, response_time (ms): %.1f"
                % (percentile, acc, d / 1000.0)
            )
        accuracy.append((percentile, acc, d / 1000.0))
    return accuracy


def ConstructEndToEndTraces(
        pred_assignments_by_process, true_assignments_by_process, in_spans_by_process, all_spans
):
    def OrderTraces(end_to_end_traces):
        for trace_id in end_to_end_traces:
            end_to_end_traces[trace_id].sort(
                key=lambda x: float('inf') if x is None else x.start_mus
            )

    processes = true_assignments_by_process.keys()
    processes = list(set(processes) - set(config.otel_block_list))

    true_traces = {}
    pred_traces = {}
    for process in processes:
        for in_span in in_spans_by_process[process]:
            if in_span.trace_id not in pred_traces:
                true_traces[in_span.trace_id] = []
                pred_traces[in_span.trace_id] = []
            true_assignments = true_assignments_by_process[process]
            pred_assignments = pred_assignments_by_process[process]
            for ep in true_assignments.keys():
                if in_span.GetId() not in true_assignments[ep]:
                    break
                true_span = all_spans.get(true_assignments[ep][in_span.GetId()], None)
                if true_span is None:
                    continue
                if true_span not in true_traces[in_span.trace_id]:  # 去重
                    true_traces[in_span.trace_id].append(true_span)

                # options 代表并列的执行路径。
                options = pred_assignments[ep].get(in_span.GetId(), None)
                if not isinstance(options, list):
                    options = [options]
                for option in options:
                    pred_span = all_spans.get(option, None)
                    if pred_span is None:
                        continue
                    pred_traces[in_span.trace_id].append(pred_span)

    OrderTraces(true_traces)
    OrderTraces(pred_traces)

    return true_traces, pred_traces


def PrintLatency12(trace_correctness, all_spans):
    all_traces = []
    min_time = 1.0e40
    for _, span in all_spans.items():
        if span.IsRoot():
            child_spans = copy.deepcopy(span.children_spans)
            child_spans.sort(key=lambda s: all_spans[s].start_mus)
            processing_delay = all_spans[child_spans[0]].start_mus - span.start_mus
            correct = trace_correctness[span.trace_id]
            all_traces.append((span.start_mus, processing_delay, span.trace_id, correct, 1))
            min_time = min(span.start_mus, min_time)
    all_traces.sort()
    for x in all_traces:
        x = x[0] - min_time, x[1], x[2], x[3], x[4]
        # print(x)


def ExportErrorTraceToCSV(pred_traces_e2e, trace_correctness, filename):
    with open(filename, 'w') as f:
        f.write("trace_id\t"
                "trace_start_time\t"
                "root_span_id\n")
        for trace_id, correct in trace_correctness.items():
            if correct:
                continue
        for i, span in enumerate(pred_traces_e2e.get(trace_id, [])):
            if i == 0:  # 保存 root span
                f.write(f"{span.trace_id}\t{span.start_mus}\t{span.GetId()}\n")


def ExportErrorSpansToCSV(pred_assignments_by_process, true_assignments_by_process, in_spans_by_process, filename):
    with open(filename, 'w') as f:
        f.write("endpoint\t"
                "trace_id\t"
                "child_span_id\t"
                "true_parent_span_id\t"
                "pred_parent_span_id\n")
        for process in true_assignments_by_process.keys():
            true_assignments = true_assignments_by_process[process]
            pred_assignments = pred_assignments_by_process[process]
            for in_span in in_spans_by_process[process]:
                for ep in true_assignments.keys():
                    true_span = true_assignments[ep].get(in_span.GetId(), None)
                    pred_span = pred_assignments[ep].get(in_span.GetId(), None)
                    if true_span is None or pred_span is None:
                        continue
                    # 不保存 Skip
                    if pred_span == ('Skip', 'Skip'):
                        continue
                    if pred_span != true_span:
                        f.write(
                            f"{ep}\t"
                            f"{in_span.trace_id}\t"
                            f"{in_span.GetId()}\t"
                            f"{true_span}\t"
                            f"{pred_span}\n"
                        )


def _ErrorAssignmentsTable(pred_assignments_by_process, true_assignments_by_process, in_spans_by_process, filename):
    with open(filename, 'w') as f:
        headers = ["endpoint",
                   "trace_id",
                   "child_span_id",
                   "true_parent_span_id",
                   "pred_parent_span_id", ]
        table = []
        for process in true_assignments_by_process.keys():
            true_assignments = true_assignments_by_process[process]
            pred_assignments = pred_assignments_by_process[process]
            for in_span in in_spans_by_process[process]:
                for ep in true_assignments.keys():
                    if in_span.GetId() not in true_assignments[ep]:
                        continue
                    if (
                            true_assignments[ep][in_span.GetId()]
                            != pred_assignments[ep][in_span.GetId()]
                    ):
                        table.append(
                            [ep,
                             in_span.trace_id,
                             in_span.GetId(),
                             true_assignments[ep][in_span.GetId()],
                             pred_assignments[ep][in_span.GetId()], ]
                        )
        f.write(tabulate(table, headers))
